package com.zws.cucumber.adaptation.cass;

import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
import com.datastax.oss.driver.api.core.type.*;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.internal.core.type.PrimitiveType;
import com.zws.cucumber.util.colorfullog.ZwsColorfulLogger;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;

import java.util.*;
import java.util.stream.Collectors;

import static java.lang.String.format;

/**
 * author: zws
 */
public class ZwsCqlExtractor {

    private static Logger logger = ZwsColorfulLogger.of(ZwsCqlExtractor.class);

    private static char[] KEYWORD_INSERT = new char[] { 'I', 'N', 'S', 'E', 'R', 'T' };
    private static char[] KEYWORD_INTO = new char[] { 'I', 'N', 'T', 'O' };
    private static char[] KEYWORD_VALUES = new char[] { 'V', 'A', 'L', 'U', 'E', 'S' };
    private static char[] KEYWORD_UPDATE = new char[] { 'U', 'P', 'D', 'A', 'T', 'E' };
    private static char[] KEYWORD_NAN = new char[] { 'N', 'A', 'N' };
    private static char[] KEYWORD_INFINITY = new char[] { 'I', 'N', 'F', 'I', 'N', 'I', 'T', 'Y' };
    private static char[] KEYWORD_NULL = new char[] { 'N', 'U', 'L', 'L' };
    private static char[] KEYWORD_TRUE = new char[] { 'T', 'R', 'U', 'E' };
    private static char[] KEYWORD_FALSE = new char[] { 'F', 'A', 'L', 'S', 'E' };

    private static char[] KEYWORD_SELECT = new char[] { 'S', 'E', 'L', 'E', 'C', 'T' };
    private static char[] KEYWORD_FROM = new char[] { 'F', 'R', 'O', 'M' };

    private static char[] KEYWORD_DELETE = new char[] { 'D', 'E', 'L', 'E', 'T', 'E' };

    private String cql;
    private int idx;
    private CodecRegistry codecRegistry;
    private boolean isPreparedCql;

    private ZwsCqlExtractor(String cql, CodecRegistry codecRegistry, boolean isPreparedCql) {
        this.cql = cql;
        this.idx = 0;
        this.codecRegistry = codecRegistry;
        this.isPreparedCql = isPreparedCql;
    }

    private static ZwsCqlInfo extract(String cql, CodecRegistry codecRegistry, boolean isPreparedCql) {
        return new ZwsCqlExtractor(cql, codecRegistry, isPreparedCql).extract();
    }

    public static ZwsCqlInfo extract(String cql, CodecRegistry codecRegistry) {
        return extract(cql, codecRegistry, false);
    }

    public static ZwsCqlInfo extract(String cql) {
        return extract(cql, null, true);
    }

    public ZwsCqlInfo extract() {
//        logger.trace("try to extract ZwsCqlInfo from cql: {}", cql);
        if (StringUtils.isBlank(cql)) {
            return null;
        }
        // skip whitespace
        skipWhitespaces();
        // insert or update
        if (lookaheadIgnoreCase('I')) {
            return extractInsertCql();
        } else if (lookaheadIgnoreCase('U') && lookaheadKeyword(KEYWORD_UPDATE)) {
            return extractUpdateCql();
        } else if (lookaheadIgnoreCase('S')) {
            return extractSelectCql();
        } else if (lookaheadIgnoreCase('D') && lookaheadKeyword(KEYWORD_DELETE)) {
            return extractDeleteCql();
        } else {
            return null;
        }
    }

    private boolean move() {
        idx++;
        return true;
    }

    private boolean lookahead(char expected) {
        if (idx < cql.length()) {
            char actual = cql.charAt(idx);
            return expected == actual;
        }
        return false;
    }

    private boolean lookahead(char expected, int idx) {
        if (idx < cql.length()) {
            char actual = cql.charAt(idx);
            return expected == actual;
        }
        return false;
    }

    private boolean lookaheadIgnoreCase(char expected) {
        if (idx < cql.length()) {
            char actual = cql.charAt(idx);
            if (expected >= 'a' && expected <= 'z') {
                expected -= 32;
            }
            if (actual >= 'a' && actual <= 'z') {
                actual -= 32;
            }
            return expected == actual;
        }
        return false;
    }

    private boolean lookaheadIgnoreCase(char expected, int idx) {
        if (idx < cql.length()) {
            char actual = cql.charAt(idx);
            if (expected >= 'a' && expected <= 'z') {
                expected -= 32;
            }
            if (actual >= 'a' && actual <= 'z') {
                actual -= 32;
            }
            return expected == actual;
        }
        return false;
    }

    private boolean lookaheadIdentifier() {
        if (idx < cql.length()) {
            char actual = cql.charAt(idx);
            return Character.isLetter(actual) || actual == '"';
        }
        return false;
    }

    private boolean lookaheadKeyword(char[] keyword) {
        int i = 0;
        while (i < keyword.length && lookaheadIgnoreCase(keyword[i], idx + i)) {
            i++;
        }
        return i == keyword.length;
    }

    private void matchKeyword(char[] keyword) {
        int i = 0;
        for (; i < keyword.length && idx < cql.length(); i++) {
            matchIgnoreCase(keyword[i]);
        }
        if (i < keyword.length) {
            throw new IllegalStateException();
        }
    }

    private void match(char expected) {
        char actual = cql.charAt(idx);
        if (expected != actual) {
            throw new IllegalStateException(format("match exactly, expected is '%s', but actual is '%s'\n%s",
                    expected, actual, buildMatchErrorMessage()));
        }
        idx++;
    }

    private void matchIgnoreCase(char expected) {
        char actual = cql.charAt(idx);
        if (expected >= 'a' && expected <= 'z') {
            expected -= 32;
        }
        if (actual >= 'a') {
            actual -= 32;
        }
        if (expected != actual) {
            throw new IllegalStateException(format("match ignore case, expected is '%s', but actual is '%s'\n%s",
                    expected, actual, buildMatchErrorMessage()));
        }
        idx++;
    }

    private String buildMatchErrorMessage() {
        StringBuilder sb = new StringBuilder(cql);
        sb.append("\n");
        char[] paddings = new char[idx];
        Arrays.fill(paddings, ' ');
        sb.append(paddings);
        sb.append('^');
        return sb.toString();
    }

    private void skipWhitespaces() {
        while (Character.isWhitespace(cql.charAt(idx))) { idx++; }
    }

    private void skipUntilKeyword(char[] keyword) {
        do {
            if (lookaheadKeyword(keyword)) {
                return;
            }
            idx++;
        } while (idx < cql.length());
    }

    private String extractTableName() {
        StringBuilder sb = new StringBuilder(64);
        boolean hasQuotation = false;
        if ('"' == cql.charAt(idx)) {
            match('"');
            hasQuotation = true;
        }
        if (Character.isLetter(cql.charAt(idx))) {
            do {
                sb.append(cql.charAt(idx));
            } while (++idx < cql.length() && (Character.isLetterOrDigit(cql.charAt(idx)) || '_' == cql.charAt(idx)));
        }
        if (hasQuotation) {
            match('"');
        }
        if (idx < cql.length() && '.' == cql.charAt(idx)) {
            sb.delete(0, sb.length());
            hasQuotation = false;
            match('.');
            if ('"' == cql.charAt(idx)) {
                match('"');
                hasQuotation = true;
            }
            if (Character.isLetter(cql.charAt(idx))) {
                do {
                    sb.append(cql.charAt(idx));
                } while (++idx < cql.length() && (Character.isLetterOrDigit(cql.charAt(idx)) || '_' == cql.charAt(idx)));
            }
            if (hasQuotation) {
                match('"');
            }
        }
        if (sb.length() > 0) {
            String table = sb.toString();
            logger.debug("extract table: {}", table);
            return hasQuotation ? sb.toString() : sb.toString().toLowerCase();
        } else {
            logger.debug("fail to extract table");
            return null;
        }
    }

    private String extractIdentifier() {
        StringBuilder sb = new StringBuilder(64);
        if ('"' == cql.charAt(idx)) {
            match('"');
            while ('"' != cql.charAt(idx)
                    || (idx + 1 < cql.length() && '"' == cql.charAt(idx + 1))) {

                if ('"' == cql.charAt(idx)) {
                    idx++;
                }
                sb.append(cql.charAt(idx++));
            }
            match('"');
            return sb.toString();
        } else {
            if (Character.isLetter(cql.charAt(idx))) {
                do {
                    sb.append(cql.charAt(idx));
                } while (Character.isLetterOrDigit(cql.charAt(++idx)) || '_' == cql.charAt(idx));
            }
            return sb.toString().toLowerCase();
        }
    }

    private Object extractTermValueByColumn(TableMetadata tableMetadata, String column) {
        ColumnMetadata columnMetadata = tableMetadata.getColumn(column)
                .orElseThrow(() -> new IllegalArgumentException(format("Unknown column '%s' for table '%s'", column, tableMetadata.getName().asInternal())));

        DataType dataType = columnMetadata.getType();
        return extractTermValue(dataType);
    }

    private Object extractTermValue(DataType type) {
        boolean isIdentifier = lookaheadIdentifier();
        String termLiteral = extractTermLiteral(type);
        if (type instanceof PrimitiveType && !isIdentifier) {
            return codecRegistry.codecFor(type).parse(termLiteral);
        }
        // just support primitive type, or else return literal
        return termLiteral;
    }

    private String extractTermLiteral(DataType type) {
        // null
        if (lookaheadIgnoreCase('N')
                && lookaheadIgnoreCase('U', idx + 1)
                && lookaheadIgnoreCase('L', idx + 2)
                && lookaheadIgnoreCase('L', idx + 3)) {
            matchKeyword(KEYWORD_NULL);
            return null;
        }
        // true
        if (lookaheadKeyword(KEYWORD_TRUE)) {
            matchKeyword(KEYWORD_TRUE);
            return "true";
        }
        // false
        if (lookaheadKeyword(KEYWORD_FALSE)) {
            matchKeyword(KEYWORD_FALSE);
            return "false";
        }
        // function call
        if (lookaheadIdentifier()) {
            return extractFunctionCallLiteral();
        }
        // constant
        if (type instanceof PrimitiveType) {
            switch (type.asCql(false, false)) {
                case "tinyint":
                case "smallint":
                case "int":
                case "bigint":
                case "counter":
                case "varint":
                    return extractIntegerLiteral();
                case "float":
                case "double":
                case "decimal":
                    return extractFloatLiteral();
                case "ascii":
                case "text":
                case "varchar":
                case "inet":
                    return extractStringLiteral();
                case "date":
                case "time":
                case "timestamp":
                    if (lookahead('\'')) {
                        return extractStringLiteral();
                    } else {
                        return extractIntegerLiteral();
                    }
                default:
                    throw new UnsupportedOperationException("unsupported primitive type: " + type.asCql(false, false));
            }
        }
        // collection
        if (type instanceof ListType) {
            return extractListLiteral((ListType) type);
        }
        if (type instanceof SetType) {
            return extractSetLiteral((SetType) type);
        }
        if (type instanceof MapType) {
            return extractMapLiteral((MapType) type);
        }
        // udt
        if (type instanceof UserDefinedType) {
            return extractUdtLiteral((UserDefinedType) type);
        }
        throw new UnsupportedOperationException();
    }

    private String extractStringLiteral() {
        int startIdx = idx;
        match('\'');
        while (!lookahead('\'')
                || (lookahead('\'') && lookahead('\'', idx + 1) && move())) {
            idx++;
        }
        match('\'');
        return cql.substring(startIdx, idx);
    }

    private String extractIntegerLiteral() {
        int startIdx = idx;
        if (lookahead('-')) {
            match('-');
        }
        while (Character.isDigit(cql.charAt(idx))) {
            idx++;
        }
        return cql.substring(startIdx, idx);
    }

    private String extractFloatLiteral() {
        int startIdx = idx;
        if (lookaheadIgnoreCase('N')) {
            matchKeyword(KEYWORD_NAN);
            return "NAN";

        } else if (lookaheadIgnoreCase('I')) {
            matchKeyword(KEYWORD_INFINITY);
            return "INFINITY";

        } else {
            if (lookahead('-')) {
                match('-');
            }
            while (Character.isDigit(cql.charAt(idx))) {
                idx++;
            }
            if (lookahead('.')) {
                match('.');
                while (Character.isDigit(cql.charAt(idx))) {
                    idx++;
                }
            }
            if (lookaheadIgnoreCase('E')) {
                matchIgnoreCase('E');
                if (lookahead('+')) {
                    match('+');
                } else if (lookahead('-')) {
                    match('-');
                }
                while (Character.isDigit(cql.charAt(idx))) {
                    idx++;
                }
            }
            return cql.substring(startIdx, idx);
        }
    }

    private String extractListLiteral(ListType type) {
        int startIdx = idx;
        match('[');
        skipWhitespaces();
        if (!lookahead(']')) {
            do {
                if (lookahead(',')) {
                    match(',');
                }
                skipWhitespaces();
                extractTermLiteral(type.getElementType());
                skipWhitespaces();
            } while (lookahead(','));
        }
        match(']');
        return cql.substring(startIdx, idx);
    }

    private String extractSetLiteral(SetType type) {
        int startIdx = idx;
        match('{');
        skipWhitespaces();
        if (!lookahead('}')) {
            do {
                if (lookahead(',')) {
                    match(',');
                }
                skipWhitespaces();
                extractTermLiteral(type.getElementType());
                skipWhitespaces();
            } while (lookahead(','));
        }
        match('}');
        return cql.substring(startIdx, idx);
    }

    private String extractMapLiteral(MapType type) {
        int startIdx = idx;
        match('{');
        skipWhitespaces();
        if (!lookahead('}')) {
            do {
                if (lookahead(',')) {
                    match(',');
                }
                skipWhitespaces();
                extractTermLiteral(type.getKeyType());
                skipWhitespaces();
                match(':');
                skipWhitespaces();
                extractTermLiteral(type.getValueType());
                skipWhitespaces();
            } while (lookahead(','));
        }
        match('}');
        return cql.substring(startIdx, idx);
    }

    private String extractUdtLiteral(UserDefinedType type) {
        int startIdx = idx;
        match('{');
        skipWhitespaces();
        if (!lookahead('}')) {
            do {
                if (lookahead(',')) {
                    match(',');
                }
                skipWhitespaces();
                String identifier = extractIdentifier();
                skipWhitespaces();
                match(':');
                skipWhitespaces();
                extractTermLiteral(type.getFieldTypes().get(type.firstIndexOf(identifier)));
                skipWhitespaces();
            } while (lookahead(','));
        }
        match('}');
        return cql.substring(startIdx, idx);
    }

    private String extractFunctionCallLiteral() {
        int startIdx = idx;
        int count = 1;
        extractIdentifier();
        skipWhitespaces();
        match('(');
        while (true) {
            char c = cql.charAt(idx);
            if (c == '(') {
                count++;
            }
            if (c == ')') {
                if (--count == 0) {
                    break;
                }
            }
            idx++;
        }
        match(')');
        return cql.substring(startIdx, idx);
    }

    private ZwsCqlInfo extractInsertCql() {
        logger.trace("try to extract ZwsCqlInfo from insert cql: {}", cql);
        //
        String tableName;
        ZwsCqlDmlType dmlType = ZwsCqlDmlType.Insert;
        // match insert
        matchKeyword(KEYWORD_INSERT);
        // skip whitespace
        skipWhitespaces();
        // match into
        matchKeyword(KEYWORD_INTO);
        // skip whitespace
        skipWhitespaces();
        // parse table name
        tableName = extractTableName();
        if (isPreparedCql) {
            return buildCqlInfo(tableName, dmlType);
        }
        /*
         * extend info
         */
        TableMetadata tableMetadata = ZwsCqlSessionProxyHelper.getTableMetadata(tableName);
        List<String> keyList = new ArrayList<>();
        List<Object> valList = new ArrayList<>();
        // skip whitespace
        skipWhitespaces();
        // columns
        match('(');
        do {
            skipWhitespaces();
            keyList.add(extractIdentifier());
            skipWhitespaces();
            if (lookahead(',')) {
                match(',');
                skipWhitespaces();
            }
        } while(lookaheadIdentifier());
        match(')');
        skipWhitespaces();
        // match keyword values
        matchKeyword(KEYWORD_VALUES);
        skipWhitespaces();
        // values
        match('(');
        skipWhitespaces();
        for (int i = 0; i < keyList.size(); i++) {
            if (i != 0) {
                match(',');
                skipWhitespaces();
            }
            valList.add(extractTermValueByColumn(tableMetadata, keyList.get(i)));
            skipWhitespaces();
        }
        match(')');
        // print some debug info, need to be removed
        logger.trace("partition key list: {}", tableMetadata.getPartitionKey().stream().map(c -> c.getName().asInternal()).collect(Collectors.toList()));
        logger.trace("key list: {}", keyList);
        logger.trace("val list: {}", valList);
        // filter partition keys
        Map<String, Object> partitionKeyAndValueMap = new HashMap<>();
        List<ColumnMetadata> partitionKeyList = tableMetadata.getPartitionKey();

        for (ColumnMetadata columnMetadata : partitionKeyList) {
            String partitionKeyColumn = columnMetadata.getName().asInternal();
            boolean found = false;
            for (int i = 0; i < keyList.size(); i++) {
                String column = keyList.get(i);
                if (Objects.equals(partitionKeyColumn, column)) {
                    partitionKeyAndValueMap.put(column, valList.get(i));
                    found = true;
                    break;
                }
            }
            if (!found) {
                throw new IllegalStateException("not found partition key column: " + partitionKeyColumn);
            }
        }

        return buildCqlExtendedInfo(tableName, dmlType, partitionKeyAndValueMap);
    }

    private ZwsCqlInfo extractSelectCql() {
        logger.trace("try to extract ZwsCqlInfo from select cql: {}", cql);
        //
        String tableName;
        ZwsCqlDmlType dmlType = ZwsCqlDmlType.Select;
        // match insert
        matchKeyword(KEYWORD_SELECT);
        // skip until from
        skipUntilKeyword(KEYWORD_FROM);
        // match
        matchKeyword(KEYWORD_FROM);
        //
        skipWhitespaces();
        tableName = extractTableName();
        return buildCqlInfo(tableName, dmlType);
    }

    private ZwsCqlInfo extractDeleteCql() {
        logger.trace("try to extract ZwsCqlInfo from delete cql: {}", cql);
        //
        String tableName;
        ZwsCqlDmlType dmlType = ZwsCqlDmlType.Delete;
        // match insert
        matchKeyword(KEYWORD_DELETE);
        skipWhitespaces();
        // skip from
        matchKeyword(KEYWORD_FROM);
        //
        skipWhitespaces();
        tableName = extractTableName();
        return buildCqlInfo(tableName, dmlType);
    }

    private ZwsCqlInfo extractUpdateCql() {
        logger.trace("try to extract ZwsCqlInfo from update cql: {}", cql);
        //
        String tableName;
        ZwsCqlDmlType dmlType = ZwsCqlDmlType.Update;
        // match insert
        matchKeyword(KEYWORD_UPDATE);
        skipWhitespaces();

        tableName = extractTableName();
        return buildCqlInfo(tableName, dmlType);
    }

    private ZwsCqlInfo buildCqlInfo(String tableName, ZwsCqlDmlType dmlType) {
        ZwsCqlInfo info = new ZwsCqlInfo();
        info.setTableName(tableName);
        info.setDmlType(dmlType);
        return info;
    }

    private ZwsCqlInfo buildCqlExtendedInfo(String tableName, ZwsCqlDmlType dmlType, Map<String, Object> partitionKeyAndValueMap) {
        ZwsCqlExtendedInfo info = new ZwsCqlExtendedInfo();
        info.setTableName(tableName);
        info.setDmlType(dmlType);
        info.setPartitionKeyAndValueMap(partitionKeyAndValueMap);
        return info;
    }

}
